import json
from channels.generic.websocket import AsyncWebsocketConsumer
import random
from asgiref.sync import sync_to_async


class ChatConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.room_name = None
        self.room_group_name = None
        self.username = None
        self.color = None

    async def connect(self):
        self.room_name = self.scope['url_route']['kwargs']['room_name']
        self.room_group_name = f'chat_{self.room_name}'

        # 从查询参数获取用户名
        query_string = self.scope.get('query_string', b'').decode()
        query_params = dict(param.split('=') for param in query_string.split('&') if '=' in param)
        self.username = query_params.get('username', 'Anonymous')

        self.color = self.generate_color()

        # 加入房间组
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )

        await self.accept()

        # 发送用户加入通知
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'chat_message',
                'message': f"{self.username} 加入了聊天室",
                'username': '系统',
                'color': '#999999'
            }
        )

        # 发送当前用户数
        await self.send_user_count()

    def generate_color(self):
        colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A', '#98D8C8']
        return random.choice(colors)

    async def disconnect(self, close_code):
        # 离开房间组
        await self.channel_layer.group_discard(
            self.room_group_name,
            self.channel_name
        )

        # 发送用户离开通知
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'chat_message',
                'message': f"{self.username} 离开了聊天室",
                'username': '系统',
                'color': '#999999'
            }
        )

        # 更新用户数
        await self.send_user_count()

    async def send_user_count(self):
        # 获取当前房间用户数
        # 注意：这是一种简化方法，实际应用中可能需要更精确的计数方式
        # 这里我们发送一个固定值1，或者可以实现更复杂的计数逻辑
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'user_count',
                'count': 1  # 这里应该替换为实际的用户计数
            }
        )

    async def receive(self, text_data):
        text_data_json = json.loads(text_data)
        message = text_data_json['message']

        # 发送消息到房间组
        await self.channel_layer.group_send(
            self.room_group_name,
            {
                'type': 'chat_message',
                'message': message,
                'username': self.username,
                'color': self.color
            }
        )

    async def chat_message(self, event):
        await self.send(text_data=json.dumps({
            'type': 'message',
            'message': event['message'],
            'username': event['username'],
            'color': event['color']
        }))

    async def user_count(self, event):
        await self.send(text_data=json.dumps({
            'type': 'user_count',
            'count': event['count']
        }))